import abc
from torch.utils.data import Dataset, DataLoader
import numpy as np
import torch
import json
from utils.label_transform import real2label
from utils.image_process import norm_added_image, norm_image

# 导入xy坐标
coord = np.load("./data/coord/normed_coord_xy.npy")


# 定义dataloader导入数据
def default_loader(img_path):
    img_numpy = np.load(img_path)
    # 将img像素点在O_source下的坐标添加到img当中
    # img_numpy = add_o_source_coord(img_numpy)
    # 对已经加入的带坐标的图像在每一个通道做归一化
    # img_tensor = norm_added_image(img_numpy)
    # 对图像本身做归一化
    img_numpy = norm_image(img_numpy.astype(np.float64)) / 255
    # 添加xy坐标
    img_numpy = np.vstack((img_numpy[np.newaxis, :], coord))
    img_tensor = torch.from_numpy(img_numpy).float().cuda()
    return img_tensor


# 定义获取本次实验所用的DRR的数据库
class Trainset(Dataset):
    def __init__(self, file_train, number_train, loader=default_loader):
        # 定义好 image 的路径
        self.images = file_train
        self.target = number_train
        self.loader = loader

    def __getitem__(self, index):
        fn = self.images[index]
        img = self.loader(fn)
        # print(self.target[index])
        target = np.array([self.target[index]["angle_x"],
                           self.target[index]["angle_y"],
                           self.target[index]["angle_z"],
                           self.target[index]["tx"],
                           self.target[index]["ty"],
                           self.target[index]["tz"]])
        target = torch.from_numpy(target).float().cuda()
        return img, target

    def __len__(self):
        return len(self.images)


# 通过json获取图片的信息
def get_info_from_json(DRR_path, fn_path):
    """

    :param DRR_path: 输入与DRR图片相关的文件路径
    :param fn_path: 输入与DRR图片信息相关的文件路径
    :return:列表形式的图片路径和信息
    """
    # 载入图片数据
    img_path = []
    img_index = []
    # 读取json文件
    with open(fn_path, 'r') as img_json:
        img_info = json.load(img_json)
        # print(jsoimgn.dumps(train_info, indent=4))
    for key in img_info:
        # print(train_info[key])
        path = DRR_path + "/" + img_info[key]['name']
        img_path.append(path)
        img_index.append(img_info[key])
    return img_path, img_index


# 实例化某一个数据库，包括图片数据的导入和将数据库封装为一个迭代器
class DRR_dataset:
    def __init__(self, DRR_path, fn_path, Batch_size):
        train_path, train_index = get_info_from_json(DRR_path=DRR_path, fn_path=fn_path)
        self.train_data = Trainset(train_path, train_index)
        self.Batch_size = Batch_size
        self.train_loader = DataLoader(self.train_data, batch_size=self.Batch_size, shuffle=True, drop_last=True)
        self.train_iter = iter(self.train_loader)

    def get_data(self):
        try:
            batch_x, batch_y = next(self.train_iter)
        except StopIteration:
            train_loader = DataLoader(self.train_data, batch_size=self.Batch_size, shuffle=True, drop_last=True)
            train_iter = iter(train_loader)
            batch_x, batch_y = next(train_iter)

        return batch_x, batch_y
